# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

"""
异步函数代码生成测试
"""

from tests.pytests.helpers.codegen_helpers import compile_source, assert_code_contains


def test_async_await_result_in_co_yield_block():
    """测试 await 结果变量在 co_yield 块中的生成

    关键验证点：
    1. v1 是成员变量（跨 await 使用）
    2. v1 的赋值应该在下一个 co_yield 块内，而不是独立一行
    3. 符合 async_ruler.md 的情况3规则
    """
    src = """
async def async_step(x: int) -> int:
    return x * 2

async def compute() -> int:
    a = 0
    b = 0
    for i in range(3):
        a = a + i
    v1 = await async_step(a)
    
    if v1 > 5:
        b = 10
    else:
        b = 20
    
    c = 0
    for i in range(2):
        for j in range(2):
            c = c + i + j
    v2 = await async_step(c)
    
    if v2 > 10:
        result = v1 + v2 + b
    else:
        result = v1 + v2 - b
    
    return result
"""
    hpp, cpp = compile_source(src)

    # 验证完整的 execute 函数
    assert_code_contains(
        cpp,
        """
        mc::result<int> ${w+}_compute::impl_type::execute(${w+}_compute::impl_type::resume_value_type& __resume_value, mcpy::async_context& __ctx) {
            co_reenter(this) {
                co_yield {
                    int a = 0;
                    b     = 0;
                    for (int i = 0; i < 3; i += 1) {
                        a = a + i;
                    }
                    return await(${w+}_async_step{}(a), __ctx);
                }
                co_yield {
                    v1 = __resume_value.get<int>();
                    if (v1 > 5) {
                        b = 10;
                    } else {
                        b = 20;
                    }
                    int c = 0;
                    for (int i = 0; i < 2; i += 1) {
                        for (int j = 0; j < 2; ++j) {
                            c = c + i + j;
                        }
                    }
                    return await(${w+}_async_step{}(c), __ctx);
                }
                int v2 = __resume_value.get<int>();
                if (v2 > 10) {
                    int result = v1 + v2 + b;
                } else {
                    result = v1 + v2 - b;
                }
                return set_result(result);
            }
            return set_result(__resume_value);
        }
        """,
        use_markers=True,
    )


def test_async_local_variable_in_co_yield():
    """测试局部变量在 co_yield 块中的生成

    所有变量都是局部变量，不跨 await 使用
    """
    src = """
async def step1() -> int:
    return 1

async def step2(x: int) -> int:
    return x

async def step3(x: int) -> int:
    return x

async def simple() -> int:
    a = await step1()
    temp = a + 1
    b = await step2(temp)
    result = b * 3
    c = await step3(result)
    return c
"""
    hpp, cpp = compile_source(src)

    # 验证完整的 execute 函数
    assert_code_contains(
        cpp,
        """
        mc::result<int> ${w+}_simple::impl_type::execute(${w+}_simple::impl_type::resume_value_type& __resume_value, mcpy::async_context& __ctx) {
            co_reenter(this) {
                co_yield_return await(${w+}_step1{}(), __ctx);
                co_yield {
                    int a    = __resume_value.get<int>();
                    int temp = a + 1;
                    return await(${w+}_step2{}(temp), __ctx);
                }
                co_yield {
                    int b      = __resume_value.get<int>();
                    int result = b * 3;
                    return await(${w+}_step3{}(result), __ctx);
                }
                int c = __resume_value.get<int>();
                return set_result(c);
            }
            return set_result(__resume_value);
        }
        """,
        use_markers=True,
    )


def test_async_member_variable_without_control_flow():
    """测试成员变量但没有控制流的场景

    x 是成员变量（跨 await 使用），展示情况2和情况3的组合
    """
    src = """
async def step1() -> int:
    return 1

async def step2(x: int) -> int:
    return x

async def compute() -> int:
    x = 1
    y = await step1()
    z = x + y
    result = await step2(z)
    return result
"""
    hpp, cpp = compile_source(src)

    # 验证完整的 execute 函数
    assert_code_contains(
        cpp,
        """
        mc::result<int> ${w+}_compute::impl_type::execute(${w+}_compute::impl_type::resume_value_type& __resume_value, mcpy::async_context& __ctx) {
            co_reenter(this) {
                x = 1;
                co_yield_return await(${w+}_step1{}(), __ctx);
                co_yield {
                    int y = __resume_value.get<int>();
                    int z = x + y;
                    return await(${w+}_step2{}(z), __ctx);
                }
                int result = __resume_value.get<int>();
                return set_result(result);
            }
            return set_result(__resume_value);
        }
        """,
        use_markers=True,
    )


def test_async_no_control_flow_before_await():
    """测试 await 前没有控制流的场景

    只有简单的赋值语句，符合情况3（有局部变量定义）
    """
    src = """
async def get_value(x: int) -> int:
    return x * 2

async def process() -> int:
    a = 10
    b = 20
    v = await get_value(a + b)
    return v + 5
"""
    hpp, cpp = compile_source(src)

    # 验证完整的 execute 函数
    assert_code_contains(
        cpp,
        """
        mc::result<int> ${w+}_process::impl_type::execute(${w+}_process::impl_type::resume_value_type& __resume_value, mcpy::async_context& __ctx) {
            co_reenter(this) {
                co_yield {
                    int a = 10;
                    int b = 20;
                    return await(${w+}_get_value{}(a + b), __ctx);
                }
                int v = __resume_value.get<int>();
                return set_result(v + 5);
            }
            return set_result(__resume_value);
        }
        """,
        use_markers=True,
    )


def test_list_append_await():
    src = """
async def get_value(x: int) -> int:
    return x * 2

async def process() -> list[int]:
    result: list[int] = []
    __await_temp_0 = await get_value(1)
    result.append(__await_temp_0)
    __await_temp_1 = await get_value(2)
    result.append(__await_temp_1)
    __await_temp_2 = await get_value(3)
    result.append(__await_temp_2)

    return result
"""
    hpp, cpp = compile_source(src)
    assert_code_contains(
        cpp,
        """
        mc::result<mc::array<int>> ${w+}_process::impl_type::execute(${w+}_process::impl_type::resume_value_type& __resume_value, mcpy::async_context& __ctx) {
            co_reenter(this) {
                result = mc::array<int>{};
                co_yield_return await(${w+}_get_value{}(1), __ctx);
                co_yield {
                    int __await_temp_0 = __resume_value.get<int>();
                    result.push_back(__await_temp_0);
                    return await(${w+}_get_value{}(2), __ctx);
                }
                co_yield {
                    int __await_temp_1 = __resume_value.get<int>();
                    result.push_back(__await_temp_1);
                    return await(${w+}_get_value{}(3), __ctx);
                }
                int __await_temp_2 = __resume_value.get<int>();
                result.push_back(__await_temp_2);
                return set_result(result);
            }
            return set_result(__resume_value);
        }
        """,
        use_markers=True,
    )


def test_async_void_return_type():
    """测试无返回值的异步函数（返回 None）

    验证点：
    1. 返回类型应该是 mc::result<void>
    2. co_reenter 块末尾应该调用 set_result()
    """
    src = """
async def sleep_ms(duration_ms: int):
    await sleep(duration_ms)

async def process() -> int:
    await sleep(10)
    await sleep_ms(10)
    return 10
"""
    hpp, cpp = compile_source(src)
    assert_code_contains(
        cpp,
        """
        mc::result<void> ${w+}_sleep_ms::impl_type::execute(${w+}_sleep_ms::impl_type::resume_value_type& __resume_value, mcpy::async_context& __ctx) {
            co_reenter(this) {
                co_yield_return await(sleep(duration_ms), __ctx);
            }
            return set_result(__resume_value);
        }
        """,
        use_markers=True,
    )


def test_async_basic_try_except():
    """测试基本的 try/except 异常处理

    验证点：
    1. try 块内的 await 应该使用 co_yield try 形式
    2. await 结果获取在 co_yield try 块之后
    3. 异常处理标签应该在 await 结果获取之后
    4. except 块应该正确处理异常
    """
    src = """
async def async_func_that_throws() -> int:
    raise ValueError("test error")
    return 42

async def test_basic_try_except():
    result = 0
    try:
        result = await async_func_that_throws()
    except ValueError:
        result = -1
    return result
"""
    hpp, cpp = compile_source(src)

    # 验证完整的 execute 函数（使用新的宏格式）
    assert_code_contains(
        cpp,
        """
        mc::result<int> ${w+}_test_basic_try_except::impl_type::execute(${w+}_test_basic_try_except::impl_type::resume_value_type& __resume_value, mcpy::async_context& __ctx) {
            std::exception_ptr __exception_ptr_try_0 = nullptr;
            co_reenter(this) {
                result = 0;
                co_yield try {
                    return await(${w+}_async_func_that_throws{}(), __ctx);
                }
                co_catch(0, CO_CLEAR_STATE);
                co_check_exception(0);
                result = __resume_value.get<int>();
                co_except(0) {
                    co_rethrow(0) catch (const mcpy::value_error&) {
                        __exception_ptr_try_0 = nullptr;
                        result = -1;
                    }
                }
                return set_result(result);
            }
            return set_result(__resume_value);
        }
        """,
        use_markers=True,
    )
